import tensorflow as tf

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("input_data/", one_hot=False)

# 定义超参数
learning_rate = 0.01
epochs = 3000
batch_size = 128
input_size = 784
class_num = 10
dropout = 0.25


# 创建神经网络
def conv_net(x_dict, n_classes, dropout, reuse, is_training):
    with tf.variable_scope('ConvNet',reuse=reuse):
        x = x_dict['images']

        # MNIST 数据集输入为784像素
        # 重塑形状为固定格式 [Height x Width x Channel]
        # 修改为张量形式 4-D: [Batch Size, Height, Width, Channel]
        x = tf.reshape(x, shape=[-1, 28, 28, 1])

        # 创建卷积层
        conv1 = tf.layers.conv2d(x, 32, 5, activation=tf.nn.relu)
        # 池化层卷积核为2步长为2
        conv1 = tf.layers.max_pooling2d(conv1, 2, 2)

        # 64个filter要用3长的kernel
        conv2 = tf.layers.conv2d(conv1, 64, 3, activation=tf.nn.relu)
        # 池化层下采样 用2*2的卷积核
        conv2 = tf.layers.max_pooling2d(conv2, 2, 2)

        fc1 = tf.contrib.layers.flatten(conv2)
        fc1 = tf.layers.dense(fc1, 1024)
        # 加入dropout层避免过拟合
        fc1 = tf.layers.dropout(fc1, rate=dropout, training=is_training)
        # 输出层预测
        out = tf.layers.dense(fc1, class_num)
    return out


# 定义模型功能函数 按照Estimator tensorflow的评估器模板建立
def model_func(features, labels, mode):
    # 建立神经网络
    # 因为dropout在预测和和训练时间上有所不同，所以需要建立2个计算流图共享相同的权重
    logits_train = conv_net(features, class_num, dropout, reuse=False,
                            is_training=True)
    logits_test = conv_net(features, class_num, dropout, reuse=True,
                           is_training=False)
    # 预测过程
    pred_classes = tf.argmax(logits_test,axis=1)
    pred_prob = tf.nn.softmax(logits_test)

    # 如果是预测模型则提前返回
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=pred_classes)

    # 定义损失函数和梯度下降法 labels转换类型
    loss_op = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_test, labels=tf.cast(labels, dtype=tf.int32)))
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
    train_op = optimizer.minimize(loss_op, global_step=tf.train.get_global_step())

    # 计算模型准确率
    acc_op = tf.metrics.accuracy(labels=labels, predictions=pred_classes)

    # TF Estimator模型评估器需要返回一个专有的类型数据包含了不同操作的数据
    estim_specs = tf.estimator.EstimatorSpec(
        mode=mode,
        predictions=pred_classes,
        loss=loss_op,
        train_op=train_op,
        eval_metric_ops={'accuracy': acc_op})

    return estim_specs


# 建立评估器
model = tf.estimator.Estimator(model_func)
# 定义训练函数所需要的输入
input_func = tf.estimator.inputs.numpy_input_fn(
    x={'images': mnist.train.images}, y=mnist.train.labels,
    batch_size=batch_size, num_epochs=None, shuffle=True)

# 训练模型
model.train(input_func, steps=epochs)

input_fn = tf.estimator.inputs.numpy_input_fn(
    x={'images': mnist.test.images}, y=mnist.test.labels,
    batch_size=batch_size, shuffle=False)
# 使用模型评估器评估
e = model.evaluate(input_fn)

print("模型的准确率为:", e['accuracy'])
